#include "projection_correct.h"

using namespace std;
using namespace cv;

void ProjectionCorrect::getProjectionMatrix(const cv::Mat& image, cv::Mat& matrix)
{
	if (image.empty()) {
		return;
	}

	vector<Point2f> cornerPoints;
	bool found = findCorners(image, cornerPoints);
	if (!found) {
		return;
	}

	vector<Vec3f> vectorPoints;
	extendPoints2Vector(cornerPoints, vectorPoints);

	Mat matrixHP = Mat::eye(3, 3, CV_32F);
	Mat matrixHA = Mat::eye(3, 3, CV_32F);
	Mat matrixHS = Mat::eye(3, 3, CV_32F);

	getProjectionHP(vectorPoints, matrixHP);
	getAffineHA(vectorPoints, matrixHP, matrixHA);
	getSimilarityHS(vectorPoints, matrixHP, matrixHA, matrixHS);

	matrix = matrixHS * matrixHA * matrixHP;
}


void ProjectionCorrect::extendPoints2Vector(const std::vector<cv::Point2f>& points, std::vector<cv::Vec3f>& vectors)
{
	if (points.empty()) {
		return;
	}

	int n = static_cast<int>(points.size());
	vectors.resize(n, Vec3f(1, 1, 1));

	for (int i = 0; i < n; ++i) {
		vectors[i].val[0] = points[i].x;
		vectors[i].val[1] = points[i].y;
	}
}


void ProjectionCorrect::getProjectionHP(const std::vector<cv::Vec3f>& points, cv::Mat& matrixHP)
{
	if (points.size() < 9) {
		return;
	}

	Vec3f line1 = crossMultiple(points[0], points[2]);
	Vec3f line2 = crossMultiple(points[6], points[8]);
	Vec3f line3 = crossMultiple(points[0], points[6]);
	Vec3f line4 = crossMultiple(points[2], points[8]);

	Vec3f cross1 = crossMultiple(line1, line2);
	Vec3f cross2 = crossMultiple(line3, line4);

	Vec3f lineInf = crossMultiple(cross1, cross2);

	const int n = 3;
	matrixHP = Mat::eye(n, n, CV_32F);
	for (int i = 0; i < n; ++i) {
		matrixHP.at<float>(n - 1, i) = lineInf.val[i];
	}
}


void ProjectionCorrect::getAffineHA(const std::vector<cv::Vec3f>& points, const cv::Mat& matrixHP, cv::Mat& matrixHA)
{
	if (points.size() < 9) {
		return;
	}

	vector<Vec3f> affinePoints(points.size());
	for (int i = 0; i < points.size(); ++i) {
		affinePoints[i] = convertPoint(matrixHP, points[i]);
	}

	vector<Vec3f> lines(4);
	lines[0] = crossMultiple(affinePoints[0], affinePoints[2]);
	lines[1] = crossMultiple(affinePoints[0], affinePoints[6]);
	lines[2] = crossMultiple(affinePoints[0], affinePoints[8]);
	lines[3] = crossMultiple(affinePoints[2], affinePoints[6]);
	
	Mat matrixS;
	getAffineS(lines, matrixS);

	Mat matrixK;
	cholesky(matrixS, matrixK);

	matrixK /= sqrtf(static_cast<float>(determinant(matrixK)));
	matrixK = matrixK.inv();

	matrixHA = Mat::eye(3, 3, CV_32F);
	Mat dest = matrixHA(Range(0, 2), Range(0, 2));
	matrixK.copyTo(dest);
}


void ProjectionCorrect::getAffineS(std::vector<cv::Vec3f>& lines, cv::Mat& S)
{
	S = Mat::eye(2, 2, CV_32F);
	if (lines.size() != 4) {
		return;
	}

	auto getParam = [](Vec3f& l, Vec3f& m) ->vector<float> {
		return { l.val[0] * m.val[0], l.val[0] * m.val[1] + l.val[1] * m.val[0], -l.val[1] * m.val[1] };
	};

	Mat A = Mat::zeros(2, 2, CV_32F);
	Mat b = Mat::zeros(2, 1, CV_32F);

	vector<float> param;

	param = getParam(lines[0], lines[1]);
	A.at<float>(0, 0) = param[0];
	A.at<float>(0, 1) = param[1];
	b.at<float>(0, 0) = param[2];

	param = getParam(lines[2], lines[3]);
	A.at<float>(1, 0) = param[0];
	A.at<float>(1, 1) = param[1];
	b.at<float>(1, 0) = param[2];

	if (determinant(A) == 0) {
		return;
	}

	Mat X = A.inv() * b;
	S.at<float>(0, 0) = X.at<float>(0, 0);
	S.at<float>(0, 1) = S.at<float>(1, 0) = X.at<float>(1, 0);
}


void ProjectionCorrect::cholesky(cv::Mat& A, cv::Mat& L)
{
	// A = L*L.T
	L = A.clone();
	if (A.empty() || A.rows != A.cols) {
		return;
	}

	Cholesky((float*)L.ptr(), L.step, L.rows, nullptr, 0, 0);
	for (int i = 0; i < L.rows - 1; ++i) {
		for (int j  = i + 1; j < L.cols; ++j) {
			L.at<float>(i, j) = 0;
		}
	}
}


void ProjectionCorrect::getSimilarityHS(const std::vector<cv::Vec3f>& points, const cv::Mat& matrixHP, const cv::Mat& matrixHA, cv::Mat& matrixHS)
{
	matrixHS = Mat::eye(3, 3, CV_32F);
	if (points.size() < 9) {
		return;
	}

	Mat H = matrixHA * matrixHP;
	vector<Vec3f> cvtVecs(points.size());
	vector<Vec2f> srcPoints(points.size()), dstPoints(points.size());
	for (int i = 0; i < cvtVecs.size(); ++i) {
		cvtVecs[i] = convertPoint(H, points[i]);
		for (int j = 0; j < 2; ++j) {
			srcPoints[i].val[j] = cvtVecs[i].val[j] / cvtVecs[i].val[2];
			dstPoints[i].val[j] = points[i].val[j] / points[i].val[2];
		}
	}

	auto moldLength = [](Vec2f& p) {
		return sqrtf(p.val[0] * p.val[0] + p.val[1] * p.val[1]);
	};
	auto getLength = [](Vec2f& p, Vec2f& q) {
		return sqrtf((p.val[0] - q.val[0]) * (p.val[0] - q.val[0]) + (p.val[1] - q.val[1]) * (p.val[1] - q.val[1]));
	};

	float srcLength =
		getLength(srcPoints[0], srcPoints[2]) +
		getLength(srcPoints[6], srcPoints[8]) +
		getLength(srcPoints[0], srcPoints[6]) +
		getLength(srcPoints[2], srcPoints[8]);

	float dstLength =
		getLength(dstPoints[0], dstPoints[2]) +
		getLength(dstPoints[6], dstPoints[8]) +
		getLength(dstPoints[0], dstPoints[6]) +
		getLength(dstPoints[2], dstPoints[8]);

	float scale = dstLength / srcLength;
	float cosTheta = srcPoints[4].dot(dstPoints[4]) / (moldLength(srcPoints[4]) * moldLength(dstPoints[4]));
	cosTheta = static_cast<float>(cosTheta > 1.0 ? 1.0 : (cosTheta < -1.0 ? -1.0 : cosTheta));
	float theta = acosf(cosTheta);

	if (srcPoints[4].val[1] * dstPoints[4].val[1] > 0) {
		theta = static_cast<float>(CV_2PI - theta);
	}

	Mat dest = matrixHS(Range(0, 2), Range(0, 2));
	dest.at<float>(0, 0) = scale * cos(theta);
	dest.at<float>(0, 1) = -scale * sin(theta);
	dest.at<float>(1, 0) = scale * sin(theta);
	dest.at<float>(1, 1) = scale * cos(theta);

	Vec3f newSrcVec = convertPoint(matrixHS, cvtVecs[4]);
	Vec2f newSrcPoint;
	newSrcPoint.val[0] = newSrcVec[0] / newSrcVec[2];
	newSrcPoint.val[1] = newSrcVec[1] / newSrcVec[2];

	float tx = dstPoints[4].val[0] - newSrcPoint.val[0];
	float ty = dstPoints[4].val[1] - newSrcPoint.val[1];

	matrixHS.at<float>(0, 2) = tx;
	matrixHS.at<float>(1, 2) = ty;
}


cv::Vec3f ProjectionCorrect::convertPoint(const cv::Mat& H, const cv::Vec3f& point)
{
	Mat result = H * point;
	Vec3f vec;
	for (int i = 0; i < 3; ++i) {
		vec.val[i] = result.at<float>(i, 0);
	}

	return vec;
}


cv::Vec3f ProjectionCorrect::crossMultiple(const cv::Vec3f& v1, const cv::Vec3f& v2)
{
	Vec3f v = v1.cross(v2);
	double length = sqrtf(v.val[0] * v.val[0] + v.val[1] * v.val[1] + v.val[2] * v.val[2]);
	return v / length;
}


